package com.hhf.rrd.usage.state;

import com.google.common.collect.Lists;
import org.apache.flink.api.common.functions.RichFlatMapFunction;
import org.apache.flink.api.common.state.MapState;
import org.apache.flink.api.common.state.MapStateDescriptor;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.KeyedStream;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.util.Collector;

import java.util.*;

/**
 * state
 *
 * @author huanghaifeng15
 * @date 2022/2/10 10:21
 **/
public class StateApp {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

        List<Tuple2<Long,Long>> list = new ArrayList<>();
        list.add(Tuple2.of(1L, 3L));
        list.add(Tuple2.of(1L, 7L));
        list.add(Tuple2.of(2L, 4L));
        list.add(Tuple2.of(1L, 5L));
        list.add(Tuple2.of(2L, 2L));
        list.add(Tuple2.of(3L, 5L));
        list.add(Tuple2.of(5L, 5L));
        list.add(Tuple2.of(7L, 5L));
        list.add(Tuple2.of(9L, 5L));

        KeyedStream<Tuple2<Long, Long>, Long> keyedStream = env.fromCollection(list)
                .keyBy(x -> x.f0);

        // 使用 MapState 求平均数
        SingleOutputStreamOperator<Tuple2<Long, Integer>> streamMapOperator = keyedStream.flatMap(new RichFlatMapFunction<Tuple2<Long, Long>, Tuple2<Long, Integer>>() {

            private transient MapState<String,Long> mapState;

            @Override
            public void open(Configuration parameters) throws Exception {
                MapStateDescriptor<String,Long> descriptor = new MapStateDescriptor<>("avg", String.class, Long.class);
                mapState = getRuntimeContext().getMapState(descriptor);
            }

            @Override
            public void flatMap(Tuple2<Long, Long> value, Collector<Tuple2<Long, Integer>> out) throws Exception {
                mapState.put(UUID.randomUUID().toString(), value.f1);

                Iterable<Long> mapStateIterable = mapState.values();
                List<Long> elementList = Lists.newArrayList(mapStateIterable);

                // 如果达到 3 条，就求平均数
                if (elementList.size() == 3){
                    int count = 0;
                    int sum = 0;

                    for (Long element : elementList) {
                        count += 1;
                        sum += element;
                    }

                    int avg = sum / count;
                    out.collect(Tuple2.of(value.f0, avg));

                    mapState.clear();
                }
            }
        });


        streamMapOperator.print();

        // 使用 ValueState 求平均数
        SingleOutputStreamOperator<Tuple2<Long, Integer>> streamValueOperator = keyedStream.flatMap(new RichFlatMapFunction<Tuple2<Long, Long>, Tuple2<Long, Integer>>() {

            private transient ValueState<Tuple2<Long,Long>> valueState;

            @Override
            public void open(Configuration parameters) throws Exception {
                ValueStateDescriptor<Tuple2<Long,Long>> descriptor = new ValueStateDescriptor<>("avg", Types.TUPLE(Types.LONG, Types.LONG));
                valueState = getRuntimeContext().getState(descriptor);
            }

            @Override
            public void flatMap(Tuple2<Long, Long> value, Collector<Tuple2<Long, Integer>> out) throws Exception {
                Tuple2<Long, Long> currentState = valueState.value();

                if (null == currentState){
                    currentState = Tuple2.of(0L,0L);
                }

                currentState.f0 += 1;
                currentState.f1 += value.f1;

                valueState.update(currentState);

                // 如果达到 3 条，就求平均数
                if (currentState.f0 == 3){
                    long avg = currentState.f1 / currentState.f0;
                    out.collect(Tuple2.of(value.f0, Integer.parseInt(String.valueOf(avg))));

                    valueState.clear();
                }
            }
        });

        streamValueOperator.print();

        // 执行任务
        env.execute();
    }
}
